package org.lql.chaos.kafka.example.partition;

import lombok.extern.slf4j.Slf4j;
import org.apache.kafka.clients.producer.Partitioner;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.header.Header;
import org.apache.kafka.common.header.Headers;
import org.lql.chaos.kafka.example.context.PartitionStrategyContext;

import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.Map;

/**
 * @author: liangqinglong
 * @date: 2025-08-10 13:09
 * @description: 总控分区器
 * - 优先从 Header 获取分区策略
 * - 没有 Header 时按默认策略
 * - 策略可动态扩展
 **/
@Slf4j
public class MasterPartitioner implements Partitioner {

	private final Map<String, CustomPartitionStrategy> strategies = new HashMap<>();

	@Override
	public void configure(Map<String, ?> configs) {
		strategies.put("HASH", new HashPartitionStrategy());
		strategies.put("RANDOM", new RandomPartitionStrategy());
		strategies.put("DEFAULT", new DefaultPartitionStrategy());
		log.info("MasterPartitioner 初始化完成, 策略列表: {}", strategies.keySet());
	}

	@Override
	public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
//		String strategyName = extractStrategyFromHeaders(value);
//		if (strategyName != null && strategies.containsKey(strategyName)) {
//			log.debug("使用策略 [{}] 进行分区", strategyName);
//			return strategies.get(strategyName).partition(topic, key, keyBytes, value, valueBytes, cluster);
//		}
//		log.debug("未指定策略, 使用默认策略进行分区");
		String strategyName = PartitionStrategyContext.getStrategy();
		if (strategyName != null && strategies.containsKey(strategyName)) {
			log.debug("使用策略 [{}] 进行分区", strategyName);
			return strategies.get(strategyName).partition(topic, key, keyBytes, value, valueBytes, cluster);
		}
		log.debug("未指定策略, 使用默认策略进行分区");
		return strategies.get("DEFAULT").partition(topic, key, keyBytes, value, valueBytes, cluster);
	}

	@Override
	public void close() {
		// 如果 strategy 实现了 AutoCloseable，则调用 close（可选）
		strategies.values().forEach(s -> {
			if (s instanceof AutoCloseable) {
				try {
					((AutoCloseable) s).close();
				} catch (Exception e) {
					log.warn("close strategy error", e);
				}
			}
		});
	}

	private String extractStrategyFromHeaders(Object value) {
		if (value instanceof ProducerRecord) {
			Headers headers = ((ProducerRecord<?, ?>) value).headers();
			for (Header header : headers) {
				if ("partition-strategy".equals(header.key())) {
					return new String(header.value(), StandardCharsets.UTF_8);
				}
			}
		}
		return null;
	}
}
